import pandas as pd
from mlxtend.frequent_patterns import apriori, association_rules
from mlxtend.preprocessing import TransactionEncoder
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False


def load_and_preprocess_data(file_path):
    """加载并预处理交易数据"""
    try:
        print("开始加载数据...")
        # 读取Excel文件
        df = pd.read_excel(file_path)
        print(f"数据加载成功，共{len(df)}条交易记录")
        print(f"数据包含以下列：{list(df.columns)}")

        # 确保关键列存在
        required_columns = ['is_damage', 'weather_id','supplier_id', 'product_category_id']
        for col in required_columns:
            if col not in df.columns:
                raise ValueError(f"数据中缺少必要的列 '{col}'，分析无法进行")

        # 检查is_damage列是否为布尔值或0/1
        if not pd.api.types.is_bool_dtype(df['is_damage']) and not all(x in [0, 1, True, False] for x in df['is_damage'].unique()):
            print("警告：is_damage列不是布尔值或0/1类型，将尝试转换")
            df['is_damage'] = df['is_damage'].astype(bool)

        # 将数据转换为适合Apriori算法的交易格式
        print("正在准备交易数据...")
        transactions = []

        # 使用tqdm显示进度
        for _, row in tqdm(df.iterrows(), total=len(df), desc="转换数据"):
            transaction = []

            # 添加商品损坏状态
            transaction.append(f"损坏={row['is_damage']}")

            # 添加天气信息
            transaction.append(f"天气={row['weather_id']}")

            # 添加供应商信息
            transaction.append(f"供应商={row['supplier_id']}")

            # 添加产品类别信息
            transaction.append(f"类别={row['product_category_id']}")

            transactions.append(transaction)

        print(f"交易数据准备完成，共{len(transactions)}条交易")

        return transactions, df
    except Exception as e:
        print(f"数据加载失败: {e}")
        raise


def perform_apriori(transactions, min_support=0.05, min_confidence=0.7):
    """执行Apriori算法挖掘频繁项集和关联规则"""
    # 转换交易数据为适合Apriori算法的格式
    print("正在编码交易数据...")
    te = TransactionEncoder()
    te_ary = te.fit(transactions).transform(transactions)
    df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

    # 挖掘频繁项集
    print("正在挖掘频繁项集...")
    frequent_itemsets = apriori(df_encoded, min_support=min_support, use_colnames=True)

    # 生成关联规则
    print("正在生成关联规则...")
    rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=min_confidence)

    print(f"发现{len(frequent_itemsets)}个频繁项集和{len(rules)}条关联规则")
    return frequent_itemsets, rules


def analyze_damage_factors(rules):
    """分析与商品损坏相关的因素"""
    # 筛选包含"损坏=True"的规则
    damage_rules = rules[rules['consequents'].apply(lambda x: '损坏=True' in x)]

    if len(damage_rules) == 0:
        print("没有找到与商品损坏相关的规则")
        return None

    # 按置信度和提升度排序
    sorted_rules = damage_rules.sort_values(['confidence', 'lift'], ascending=False)

    print(f"\n找到{len(sorted_rules)}条与商品损坏相关的规则")
    print("Top 10 影响商品损坏的关联规则:")

    top_rules = []
    for i, (_, rule) in enumerate(sorted_rules.iterrows(), 1):
        antecedents = ', '.join(rule['antecedents'])
        confidence = rule['confidence']
        lift = rule['lift']
        support = rule['support']

        print(f"{i}. {antecedents} → 损坏=True")
        print(f"   支持度: {support:.4f}, 置信度: {confidence:.4f}, 提升度: {lift:.4f}")

        top_rules.append({
            'rule': f"{antecedents} → 损坏=True",
            'support': support,
            'confidence': confidence,
            'lift': lift
        })

        if i >= 10:
            break

    return top_rules


def visualize_factor_impact(top_rules):
    """可视化各因素对商品损坏的影响"""
    if not top_rules or len(top_rules) == 0:
        print("没有足够的数据进行可视化")
        return

    factors = [rule['rule'] for rule in top_rules]
    confidences = [rule['confidence'] for rule in top_rules]

    plt.figure(figsize=(12, 8))
    sns.barplot(x=confidences, y=factors)

    plt.title('各因素对商品损坏的影响强度（Top 10）')
    plt.xlabel('置信度')
    plt.ylabel('关联规则')
    plt.tight_layout()
    plt.savefig('factor_impact.png')
    plt.show()


def main():
    # 文件路径 - 修改为实际文件名
    file_path = 'B_xxx.xlsx'

    # 加载和预处理数据
    transactions, df = load_and_preprocess_data(file_path)

    # 执行Apriori算法
    frequent_itemsets, rules = perform_apriori(transactions)

    # 分析与商品损坏相关的因素
    top_rules = analyze_damage_factors(rules)

    # 可视化各因素对商品损坏的影响
    if top_rules:
        visualize_factor_impact(top_rules)


if __name__ == "__main__":
    main()
